import sys
import tempfile
from PyQt5.QtWidgets import QApplication, QFileSystemModel, QTreeView
from PyQt5.QtCore import QDir, QFileSystemWatcher
from zipfile import ZipFile
from io import BytesIO
import os
import shutil
import atexit
import time


class ZipTreeView(QTreeView):
    def __init__(self):
        super().__init__()

        # 创建临时文件夹
        self.temp_dir = tempfile.mkdtemp()
        # 注册退出时执行的函数
        atexit.register(self.cleanup_temp_dir)

    def start(self, zip_data):
        self.extract_zip(zip_data, self.temp_dir)

    def extract_zip(self, zip_path, extract_path):

        # 从内存中读取 zip 文件内容
        with ZipFile(zip_path, 'r') as zip_ref:
            all_entries = zip_ref.namelist()
            # print(all_entries)

            def writeFile():
                with zip_ref.open(subname) as file_data:
                    with open(file_path, 'wb') as output_file:
                        shutil.copyfileobj(file_data, output_file)

            # 解压缩 zip 文件到临时文件夹
            for subname in all_entries:
                # print(subname)
                try:
                    file_name = subname.encode('cp437').decode(
                        'gbk', 'replace')
                except Exception as e:
                    file_name = subname.encode('utf-8').decode(
                        'gbk', 'replace')
                    print(f"An error occurred: {e}")
                # print(subname)
                # file_name = file_name.replace("/", "\\", -1)
                file_path = os.path.join(extract_path, file_name)
                # print(file_path)
                if '/' or '\\' in file_name:
                    if file_name[-1] == '/' or file_name[-1] == '\\':
                        # print(1)
                        os.makedirs(file_path, exist_ok=True)
                    else:
                        # print(2)
                        folder_path = os.path.dirname(file_path)
                        os.makedirs(folder_path, exist_ok=True)
                        writeFile()
                else:
                    # print(3)
                    writeFile()

            # zip_ref.extractall(extract_path)

            # 使用 QFileSystemModel 加载解压缩后的文件结构
            self.model = QFileSystemModel()

            # 设置根路径为临时文件夹
            self.model.setRootPath(self.temp_dir)

            # 在模型中添加解压缩后的文件
            # for root, dirs, files in os.walk(self.temp_dir):
            #     for file_name in files:
            #         file_path = os.path.join(root, file_name)
            #         relative_path = os.path.relpath(
            #             file_path, self.temp_dir)
            #         self.model.mkdir(self.model.index(''), relative_path)

        # 设置 QTreeView 展示文件系统
        self.setModel(self.model)
        self.setRootIndex(self.model.index(self.temp_dir))

        # except Exception as e:
        #     print(f"An error occurred: {e}")

    def cleanup_temp_dir(self):
        try:
            abstract_model = self.model.pyqtConfigure()
            file_system_watcher = abstract_model.fileSystemWatcher()
            file_system_watcher.removePath(self.temp_dir)
            # time.sleep(1)
            # 删除临时文件夹及其内容
            shutil.rmtree(self.temp_dir)
        except Exception as e:
            print(f"An error occurred during cleanup: {e}")
            # 尝试删除失败时，打印目录内容
            print(os.listdir(self.temp_dir))


if __name__ == '__main__':
    app = QApplication(sys.argv)

    # 替换为你的 zip 文件路径或二进制数据
    # zip_data = "C:\\Users\\ZJK\\OneDrive\\桌面\\写轮眼.zip"
    # zip_data = "F:\\默认下载\\compressedFile.zip"
    # zip_data = "F:\\默认下载\\flash29banben.zip"
    # zip_data = "F:\\默认下载\\微信图片.zip"
    zip_data = "F:\\默认下载\\电脑端.zip"
    tree_view = ZipTreeView()
    tree_view.start(zip_data)
    tree_view.setWindowTitle('Zip TreeView')
    tree_view.resize(800, 600)
    tree_view.show()

    sys.exit(app.exec_())
